// Copyright (C) Kumo inc. and its affiliates.
// Author: Jeff.li lijippy@163.com
// All rights reserved.
// This program is free software: you can redistribute it and/or modify
// it under the terms of the GNU Affero General Public License as published
// by the Free Software Foundation, either version 3 of the License, or
// (at your option) any later version.
//
// This program is distributed in the hope that it will be useful,
// but WITHOUT ANY WARRANTY; without even the implied warranty of
// MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
// GNU Affero General Public License for more details.
//
// You should have received a copy of the GNU Affero General Public License
// along with this program.  If not, see <https://www.gnu.org/licenses/>.
//
#pragma once

#include <boost/random/uniform_int_distribution.hpp>
#include <string>

#include <pollux/plan/input_generator.h>
#include <pollux/vector/flat_vector.h>
#include <pollux/vector/fuzzer/generator_spec.h>

using namespace kumo::pollux::plan;

namespace kumo::pollux::exec::test {
    template<typename T>
    class ValueDistribution {
    public:
        explicit ValueDistribution(size_t size) : size_{size} {
        }

        T operator()(FuzzerGenerator &rng) const {
            // Generate inputs within [max / size_, min / size_] so that the sum of
            // size_ values won't exceed max or below min.
            T upperBound = std::numeric_limits<T>::max() / (int64_t) size_;
            T lowerBound = std::numeric_limits<T>::min() / (int64_t) size_;
            return boost::random::uniform_int_distribution<T>(
                lowerBound, upperBound)(rng);
        }

    private:
        size_t size_;
    };

    class MapUnionSumInputGenerator : public InputGenerator {
    public:
        std::vector<VectorPtr> generate(
            const std::vector<TypePtr> &types,
            VectorFuzzer &fuzzer,
            FuzzerGenerator & /*rng*/,
            memory::MemoryPool * /*pool*/) override {
            POLLUX_CHECK_EQ(types.size(), 1);
            POLLUX_CHECK(types[0]->is_map());
            auto mapType = types[0]->as_map();
            auto valueType = mapType.valueType();
            if (valueType->is_double() || valueType->is_real()) {
                return {};
            }

            std::vector<VectorPtr> result;
            auto keyType = mapType.keyType();
            auto keysVector = fuzzer.fuzzNotNull(keyType);

            VectorPtr valuesVector;
            switch (valueType->kind()) {
                case TypeKind::TINYINT:
                    valuesVector = generateValues<TypeKind::TINYINT>(valueType, fuzzer);
                    break;
                case TypeKind::SMALLINT:
                    valuesVector = generateValues<TypeKind::SMALLINT>(valueType, fuzzer);
                    break;
                case TypeKind::INTEGER:
                    valuesVector = generateValues<TypeKind::INTEGER>(valueType, fuzzer);
                    break;
                case TypeKind::BIGINT:
                    valuesVector = generateValues<TypeKind::BIGINT>(valueType, fuzzer);
                    break;
                default:
                    POLLUX_UNREACHABLE();
            }

            return {
                fuzzer.fuzzMap(
                    keysVector, valuesVector, fuzzer.getOptions().vectorSize)
            };
        }

        void reset() override {
        }

    private:
        template<TypeKind KIND>
        VectorPtr generateValues(const TypePtr &type, VectorFuzzer &fuzzer) {
            using T = typename TypeTraits<KIND>::NativeType;
            ValueDistribution<T> distribution{fuzzer.getOptions().vectorSize};
            ScalarGeneratorSpec<KIND, ValueDistribution<T> > generatorSpec{
                type, std::move(distribution), fuzzer.getOptions().nullRatio
            };

            return fuzzer.fuzz(generatorSpec);
        }
    };
} // namespace kumo::pollux::exec::test
